import aiohttp, base64
from bs4 import BeautifulSoup
from bs4.element import Tag
from Crypto.PublicKey import RSA
from Crypto.Cipher import PKCS1_v1_5
from traceback import TracebackException
from types import TracebackType

from Models.user import PortalLoginForm
from Models.result import OperationResult


class JwglPatcher:
    """
    A class used to request student's schedule from ouc's academic management system

    Usage:
    -----
    ```python
    async with JwglPatcher() as patcher:
        await patcher.get_portal_token(body) # body: PortalLoginForm
        return await patcher.get_schedule(year, term) # year: int, term: int
    ```
    """

    _pub_key = (
        "-----BEGIN RSA PUBLIC KEY-----\n"
        "MIGfMA0GCSqGSIb3DQEBAQUAA4GNADCBiQKBgQCT5zYg1mYEGoti/TrLyjv+Y/9UsqtCen6+KAW7PP81MG6Wj6N"
        "+dYtkQ8UGMq4SH2baxCSp+OtZtt3iAaUCoIR9NTRoH3deDkKxfKOUt3nJmOYgEo75C4l02R7aXxBgc6xXSnYbMRf0MlWD7"
        "+wlvVYdQ5bveqwgn3mp5QtprDwrgQIDAQAB\n"
        "-----END RSA PUBLIC KEY-----"
    )
    _cipher = PKCS1_v1_5.new(RSA.import_key(_pub_key))

    def __init__(self, session: aiohttp.ClientSession | None = None):
        self.session = session or aiohttp.ClientSession()

    async def __aenter__(self) -> "JwglPatcher":
        return self

    async def __aexit__(
        self, exc_type: Exception, exc: TracebackException, tb: TracebackType
    ):
        await self.close()

    async def close(self):
        await self.session.close()

    async def get_portal_token(self, form: PortalLoginForm) -> OperationResult:
        try:
            encoded_pwd = base64.b64encode(
                self._cipher.encrypt(form.password.encode())
            ).decode()
            self._account = form.account

            async with self.session.post(
                "https://id.ouc.edu.cn/baseLogin/postLogin",
                json={
                    "account": form.account,
                    "password": encoded_pwd,
                    "isRememberPwdOpen": 1,
                    "randomCode": "comsys-base-login",
                    "scale": 1,
                },
            ) as resp:
                if (data := await resp.json())["code"] != 200:
                    return OperationResult(is_success=False, error=data["msg"])

            async with self.session.post(
                "https://id.ouc.edu.cn/sso/login",
                params={"service": "https://my.ouc.edu.cn/cas/login"},
            ) as resp:
                if isinstance(
                    exec := BeautifulSoup(await resp.text(), "html.parser").find(
                        "input", {"name": "execution"}
                    ),
                    Tag,
                ):
                    token = exec.get("value")
                else:
                    return OperationResult(is_success=False, error="Token not found")

            await self.session.post(
                "https://id.ouc.edu.cn/sso/login",
                data={
                    "execution": token,
                    "_eventId": "submit",
                    "username": form.account,
                    "password": "00000100",
                    "service": "https://my.ouc.edu.cn/cas/login",
                },
                headers={"Content-Type": "application/x-www-form-urlencoded"},
            )

            return OperationResult(is_success=True)
        except:
            return OperationResult(is_success=False, error="Unknown error")

    async def get_schedule(self, year: int, term: int) -> OperationResult:
        await self.session.get(
            "https://id.ouc.edu.cn/sso/bridgeLogin",
            params={"username": 22040032045, "service": "http://jwgl.ouc.edu.cn/login"},
        )

        async with self.session.get(
            "http://jwgl.ouc.edu.cn/student/wsxk.xskcb.jsp",
            params={
                "params": base64.b64encode(
                    f"xn={year}&xq={term}&xh={self._account}".encode()
                ).decode()
            },
        ) as resp:
            content = (await resp.text()).replace("\n", "").replace("\t", "").replace("\r", "")

        return OperationResult(is_success=True, data=content)
